from typing import Any, Dict, Optional
from config import logger
from sqlalchemy import create_engine, text
from models import DatabaseConnection, DatabaseSchema
from datetime import datetime, date
from decimal import Decimal


class DatabaseManager:
    def __init__(self):
        self.db_connection: Optional[DatabaseConnection] = None
        self.schema = DatabaseSchema()

    async def connect_database(self, config: Dict[str, Any]):
        try:
            if self.db_connection:
                self.db_connection.close()
            # 构建数据库连接字符串，格式为mysql+pymysql://user:password@host:port/database
            connection_string = (
                f"mysql+pymysql://{config['user']}:{config['password']}@"
                f"{config['host']}:{config['port']}/{config['database']}"
            )
            engine = create_engine(
                connection_string,
                echo=False,  # 不输出SQL语句的日志
                pool_pre_ping=True,  # 启动数据库健康检查
                pool_recycle=3600,  # 连接回收的时间为3600秒
            )
            with engine.connect() as conn:
                conn.execute(text("SELECT 1"))
            self.db_connection = DatabaseConnection(
                host=config["host"],
                port=config["port"],
                user=config["user"],
                password=config["password"],
                database=config["database"],
                engine=engine,
            )
            await self.create_tables()
            return {
                "success": True,
                "message": f"成功连接到了数据库{config['database']}",
                "connection_info": {
                    "host": config["host"],
                    "port": config["port"],
                    "database": config["database"],
                },
            }

        except Exception as e:
            logger.error(f"数据库连接失败:{str(e)}")
            return {"success": False, "message": f"数据库连接失败:{str(e)}"}

    async def execute_sql(self, sql: str) -> Dict[str, Any]:
        if not self.db_connection or not self.db_connection.engine:
            return {"success": False, "message": "数据库未连接，请先连接数据库"}
        try:
            with self.db_connection.engine.connect() as connection:
                result = connection.execute(text(sql))
                if sql.strip().upper().startswith("SELECT"):
                    rows = result.fetchall()
                    columns = result.keys()
                    serializable_rows = []
                    for row in rows:
                        row_dict = {}
                        for i, col in enumerate(columns):
                            value = row[i]
                            if isinstance(value, (date, datetime)):
                                value = value.isoformat()
                            elif isinstance(value, Decimal):
                                value = float(value)
                            row_dict[col] = value
                        serializable_rows.append(row_dict)

                return {
                    "success": True,
                    "message": f"执行SQL成功，共返回条记录",
                    "data": serializable_rows,
                    "sql": sql,
                }

        except Exception as e:
            logger.error(f"SQL执行失败:{str(e)}")
            return {
                "success": False,
                "message": f"SQL执行失败:{str(e)}",
                "error": str(e),
            }

    async def create_tables(self):
        if self.db_connection and self.db_connection.engine:
            self.schema.get_metadata().create_all(self.db_connection.engine)

    def get_database_info(self) -> Dict[str, Any]:
        if self.db_connection:
            return {
                "host": self.db_connection.host,
                "port": self.db_connection.port,
                "database": self.db_connection.database,
            }
        else:
            return {"success": False, "connected": False, "message": "数据库尚未连接"}

    def close(self) -> None:
        if self.db_connection:
            self.db_connection.close()
            self.db_connection = None
